# Copyright 2024 CHINA MERCHANTS BANK CO., LTD.
# Copyright 2024 Huawei Technologies Co., Ltd
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================


"""
fastchat stanford alpaca data convert tools.
"""
import argparse
import json
import os

import pathlib


def main(data_path, output_path):
    data_path = pathlib.Path(data_path)
    with data_path.open() as f:
        data = json.load(f)

    sources = []
    for example in data:
        if example.get("input", "") == "":
            sources.append(example['instruction'])
        else:
            instruction = example['instruction']
            if instruction[-1] == ".":
                instruction = instruction[:-1]
            instruction = instruction + ": " + example['input']
            sources.append(instruction)

    targets = []
    for example in data:
        targets.append(example['output'])

    systems = []
    for example in data:
        if "system" in example:
            systems.append(example['system'])

    flags_ = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
    with os.fdopen(os.open(output_path, flags_, 0o750), 'w', encoding='utf-8') as f:
        cnt = 0
        if not systems:
            for s, t in zip(sources, targets):
                new_data = {"id": cnt, "conversations": [{"role": "user", "content": s,},{"role": "assistant", "content": t,}]}
                json_str = json.dumps(new_data, ensure_ascii=False)
                f.write(json_str+ '\n' )
                cnt += 1
        elif len(sources)==len(systems):
            for s, t, sp in zip(sources, targets, systems):
                new_data = {"id": cnt, "conversations": [{"role": "user", "content": s,},{"role": "assistant", "content": t,}], "system": sp}
                json_str = json.dumps(new_data, ensure_ascii=False)
                f.write(json_str+ '\n' )
                cnt += 1
        else:
            raise Exception("Some data samples doesn't have system prompt. Check your data format.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--alpaca_file_path", type=str, default="alpaca-data.jsonl")
    parser.add_argument(
        "--sharegpt_file_path", type=str, default="sharegpt.jsonl"
    )
    args = parser.parse_args()
    main(args.alpaca_file_path, args.sharegpt_file_path)